
#include "vaapi_decoder.hpp"

enum AVPixelFormat get_format_vaapi(AVCodecContext *ctx, const enum AVPixelFormat *pix_fmts)
{
	const enum AVPixelFormat *p;

	for (p = pix_fmts; *p != -1; p++)
	{
		if (*p == AV_PIX_FMT_VAAPI) return *p;
	}

	fprintf(stderr, "Failed to get HW surface format.\n");
	return AV_PIX_FMT_NONE;
}

vaapi_decoder::vaapi_decoder()
{
}

vaapi_decoder::~vaapi_decoder()
{
	if (this->prepared) close_stream();
}

/**
 * @brief close rtsp stream and free all data structs
 *
 * @return int always 0
 */
int vaapi_decoder::close_stream()
{
	this->prepared = 0;
	av_frame_unref(this->rgb_frame);
	av_frame_unref(this->av_frame);
	av_frame_unref(this->hw_frame);
	av_frame_unref(this->nv12_frame);
	av_packet_free(&this->packet);
	avcodec_free_context(&this->av_codec_ctx);
	avformat_close_input(&if_ctx);
	return 0;
}

/**
 * @brief open rtsp stream and prepare a decoder for the stream
 *
 * @param input_url rtsp stream url
 *
 * @return int 0 on success, <0 if failed
 */
int vaapi_decoder::open_stream(std::string input_url, std::string render_path,
							   std::string scale_size)
{
	int ret;
	std::cout << "trying to open input with tcp" << std::endl;
	AVDictionary *options = NULL;
	av_dict_set(&options, "rtsp_transport", "tcp", 0);
	av_dict_set(&options, "buffer_size", "102400", 0);
	av_dict_set(&options, "stimeout", "2000000", 0);
	av_dict_set(&options, "max_delay", "500000", 0);
	// open input stream
	ret = avformat_open_input(&this->if_ctx, input_url.c_str(), NULL, &options);
	if (ret < 0)
	{
		std::cerr << "failed to open rtsp url: " << ret << std::endl;
		return -1;
	}
	std::cout << "open input stream" << std::endl;

	// get stream info
	std::cout << "finding stream information, this may take a while..." << std::endl;
	ret = avformat_find_stream_info(this->if_ctx, NULL);
	if (ret < 0)
	{
		std::cerr << "failed to find stream info: " << ret << std::endl;
		return -1;
	}
	std::cout << "found stream info" << std::endl;

	// find a video stream
	ret = av_find_best_stream(this->if_ctx, AVMEDIA_TYPE_VIDEO, -1, -1,
							  (const AVCodec **)&this->av_codec, 0);
	if (ret < 0)
	{
		std::cerr << "failed to find a video stream from source: " << ret << std::endl;
		return -1;
	}
	this->vstream_no = ret;
	std::cout << "found video stream" << std::endl;

	// alloc codec context
	this->av_codec_ctx = avcodec_alloc_context3(this->av_codec);
	if (this->av_codec_ctx == NULL)
	{
		std::cerr << "failed to alloc codec context" << std::endl;
		return -1;
	}
	std::cout << "alloc codec context" << std::endl;

	// pass stream info
	ret = avcodec_parameters_to_context(this->av_codec_ctx,
										this->if_ctx->streams[this->vstream_no]->codecpar);
	if (ret < 0)
	{
		std::cerr << "failed to pass parameters" << std::endl;
		return -1;
	}
	std::cout << "passed parameters to decoder" << std::endl;
	// add hardware device contex
	this->av_codec_ctx->get_format = get_format_vaapi;
	ret = av_hwdevice_ctx_create(&this->av_hwdev_ctx, AV_HWDEVICE_TYPE_VAAPI,
								 render_path.c_str(), NULL, 0);
	if (ret < 0)
	{
		std::cerr << "failed to create hwdevice ctx" << std::endl;
		return -1;
	}
	this->av_codec_ctx->hw_device_ctx = av_buffer_ref(this->av_hwdev_ctx);
	std::cout << "found hardware device" << std::endl;
	// open decoder
	ret = avcodec_open2(this->av_codec_ctx, this->av_codec, NULL);
	if (ret < 0)
	{
		std::cerr << "failed to open decoder: " << ret << std::endl;
		return -1;
	}
	std::cout << "open decoder" << std::endl;

	this->packet = av_packet_alloc();
	// filters
	// VAAPI -> NV12 -> BGR24
	const AVFilter *buffer_src = avfilter_get_by_name("buffer");
	const AVFilter *buffer_sink = avfilter_get_by_name("buffersink");
	AVFilterInOut *filter_out = avfilter_inout_alloc();
	AVFilterInOut *filter_in = avfilter_inout_alloc();
	AVBufferSrcParameters *buffer_src_ctx_par = av_buffersrc_parameters_alloc();
	AVFrame *temp_frame = av_frame_alloc();
	enum AVPixelFormat pixfmts[3] = {AV_PIX_FMT_VAAPI, AV_PIX_FMT_NV12, AV_PIX_FMT_NONE};
	char filter_args[512];
	if (!buffer_src || !buffer_sink || !filter_out || !filter_in || !buffer_src_ctx_par)
	{
		std::cerr << "failed to alloc filter variables" << std::endl;
		return -1;
	}

	this->filter_graph = avfilter_graph_alloc();
	if (this->filter_graph == NULL)
	{
		std::cerr << "failed to alloc filter graph" << std::endl;
	}

	sprintf(filter_args, "video_size=%dx%d:pix_fmt=%d:time_base=%d/%d:pixel_aspect=%d/%d",
			this->av_codec_ctx->width, this->av_codec_ctx->height, AV_PIX_FMT_VAAPI,
			this->if_ctx->streams[this->vstream_no]->time_base.num,
			this->if_ctx->streams[this->vstream_no]->time_base.den,
			this->av_codec_ctx->sample_aspect_ratio.num,
			this->av_codec_ctx->sample_aspect_ratio.den);
	std::cout << filter_args << std::endl;
	ret = avfilter_graph_create_filter(&this->buffer_src_ctx, buffer_src, "in", filter_args,
									   NULL, filter_graph);
	if (ret < 0)
	{
		std::cout << "failed to create avfilter in" << std::endl;
		return -1;
	}

	while (true)
	{
		ret = av_read_frame(this->if_ctx, this->packet);
		if (ret < 0 || this->packet->stream_index != this->vstream_no) continue;
		ret = avcodec_send_packet(this->av_codec_ctx, this->packet);
		if (ret < 0) continue;
		ret = avcodec_receive_frame(this->av_codec_ctx, temp_frame);
		if (ret < 0) continue;
		if (this->av_codec_ctx->hw_frames_ctx != NULL)
		{
			buffer_src_ctx_par->hw_frames_ctx = this->av_codec_ctx->hw_frames_ctx;
			av_frame_unref(temp_frame);
			break;
		}
	}

	av_buffersrc_parameters_set(this->buffer_src_ctx, buffer_src_ctx_par);
	ret = avfilter_graph_create_filter(&this->buffer_sink_ctx, buffer_sink, "out", NULL, NULL,
									   this->filter_graph);
	if (ret < 0)
	{
		std::cerr << "failed to create avfitler out" << std::endl;
		return -1;
	}
	ret = av_opt_set_int_list(this->buffer_sink_ctx, "pix_fmts", pixfmts, AV_PIX_FMT_NONE,
							  AV_OPT_SEARCH_CHILDREN);
	if (ret < 0)
	{
		std::cerr << "failed to set pixel format" << std::endl;
		return -1;
	}
	filter_out->name = av_strdup("in");
	filter_out->filter_ctx = this->buffer_src_ctx;
	filter_out->pad_idx = 0;
	filter_out->next = NULL;

	filter_in->name = av_strdup("out");
	filter_in->filter_ctx = this->buffer_sink_ctx;
	filter_in->pad_idx = 0;
	filter_in->next = NULL;

	int scale_w = 320, scale_h = 320;
	sscanf(scale_size.c_str(), "%dx%d", &scale_w, &scale_h);
	sprintf(filter_args, "scale_vaapi=%d:%d", scale_w, scale_h);
	ret = avfilter_graph_parse_ptr(this->filter_graph, filter_args, &filter_in, &filter_out,
								   NULL);
	if (ret < 0)
	{
		std::cerr << "failed to parse filter string" << std::endl;
		return -1;
	}

	ret = avfilter_graph_config(this->filter_graph, NULL);
	if (ret < 0)
	{
		std::cerr << "failed to check filter graph validity" << std::endl;
		return -1;
	}

	av_freep(&buffer_src_ctx_par);
	avfilter_inout_free(&filter_in);
	avfilter_inout_free(&filter_out);

	this->hw_frame = av_frame_alloc();
	if (this->hw_frame == NULL)
	{
		std::cerr << "failed to alloc hw frame" << std::endl;
		return -1;
	}

	this->rgb_frame = av_frame_alloc();
	if (this->rgb_frame == NULL)
	{
		std::cerr << "failed to alloc rgb frame" << std::endl;
		return -1;
	}

	this->nv12_frame = av_frame_alloc();
	if (this->nv12_frame == NULL)
	{
		std::cerr << "failed to alloc nv12 frame" << std::endl;
		return -1;
	}

	this->av_frame = av_frame_alloc();
	if (this->av_frame == NULL)
	{
		std::cerr << "failed to alloc av frame" << std::endl;
		return -1;
	}

	ret = av_image_alloc(this->rgb_frame->data, rgb_frame->linesize, scale_w, scale_h,
						 AV_PIX_FMT_BGR24, 1);
	if (ret < 0)
	{
		std::cerr << "failed to fill rgb frame with BGR24 " << scale_size << std::endl;
		return -1;
	}

	this->prepared = 1;
	std::cout << "decoder prepared now" << std::endl;
	return 0;
}

int vaapi_decoder::decode_frame(uint8_t *data, int len, int *width, int *height, int *step)
{
	int ret;
	SwsContext *sws_ctx;
	if (this->prepared != 1) return -1;
	do
	{
		ret = avcodec_receive_frame(this->av_codec_ctx, this->av_frame);
		if (ret == 0)
		{
			do
			{
				// remaining frames in decoder
				ret = av_buffersrc_add_frame_flags(this->buffer_src_ctx, this->av_frame,
												   AV_BUFFERSRC_FLAG_KEEP_REF);
				if (ret < 0) break;
				ret = av_buffersink_get_frame(this->buffer_sink_ctx, this->hw_frame);
				if (ret < 0) break;

				ret = av_hwframe_transfer_data(this->nv12_frame, this->hw_frame, 0);
				if (ret < 0) break;
				sws_ctx = sws_getContext(this->nv12_frame->width, this->nv12_frame->height,
										 AV_PIX_FMT_NV12, this->nv12_frame->width,
										 this->nv12_frame->height, AV_PIX_FMT_BGR24,
										 SWS_BICUBIC, nullptr, nullptr, nullptr);
				if (sws_ctx == NULL) break;
				sws_scale(sws_ctx, (uint8_t const *const *)this->nv12_frame->data,
						  this->nv12_frame->linesize, 0, this->nv12_frame->height,
						  this->rgb_frame->data, this->rgb_frame->linesize);
				sws_freeContext(sws_ctx);
				// output
				if (this->nv12_frame->linesize[0] * this->nv12_frame->height > len)
				{
					std::cerr << "image size too large: linesize " << this->nv12_frame->linesize
							  << ", height " << this->nv12_frame->height << std::endl;
					av_frame_unref(this->hw_frame);
					av_frame_unref(this->nv12_frame);
					break;
				}
				*height = this->nv12_frame->height;  // image pixel height
				*width = this->nv12_frame->width;  // image pixel width
				*step = (this->nv12_frame->linesize[0] > 0 ? this->nv12_frame->linesize[0]
														   : -this->nv12_frame->linesize[0]) *
						3;
				// linesize might be greater than width, and might be negative
				memcpy(data, this->rgb_frame->data[0], *height * *step);  // image size in bytes
				ret = 0;
			} while (0);
			av_frame_unref(this->hw_frame);
			av_frame_unref(this->nv12_frame);
			av_frame_unref(this->av_frame);
			return ret;
		}
		else if (ret == AVERROR(EAGAIN))
		{
			av_packet_unref(this->packet);
			ret = av_read_frame(this->if_ctx, this->packet);
			if (ret < 0)
			{
				// eof
				std::cerr << "failed to read frame from input: " << ret << std::endl;
				return -1;
			}

			if (this->packet->stream_index == this->vstream_no)
			{
				ret = avcodec_send_packet(this->av_codec_ctx, this->packet);
				if (ret < 0)
				{
					// failed to send
					std::cerr << "failed to send packet to decoder: " << ret << std::endl;
					return -1;
				}
			}
		}
		else
		{
			std::cerr << "unexpected return, failed to receive frame: " << ret << std::endl;
			return -1;
		}
		// continue loop, try receive frame from decoder
	} while (true);
	return -1;
}

// int main(int argc, char **argv)
// {
// 	int ret, cnt = 0;
// 	if (argc < 2)
// 	{
// 		printf("Usage: %s <stream url>\n", argv[0]);
// 		return 0;
// 	}
// 	auto inf = vaapi_decoder(std::string(argv[1]));
// 	AVFrame *frame = NULL;
// 	frame = av_frame_alloc();
// 	if (frame == NULL)
// 	{
// 		std::cerr << "failed to alloc frame" << std::endl;
// 		return 0;
// 	}
// 	cnt = 0;
// 	while (1)
// 	{
// 		ret = inf.decode_frame(frame);
// 		if (ret <= 0)
// 		{
// 			std::cout << "met " << ret << std::endl;
// 			break;
// 		}
// 		std::cout << ".";
// 		cnt++;
// 	}
// 	std::cout << "decoded: " << cnt << " frames" << std::endl;
// 	return 0;
// }
